/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

import org.slf4j.Logger;
import org.sss.common.InputDrivenTask;
import org.sss.common.TaskSliceSubmitter;
import org.sss.common.io.DataReader;
import org.sss.common.io.KeyFormatter;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.Decoder;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.mapreduce.TupleGroupScaner;
import org.sss.mapreduce.datatype.Packable;

/**
 * Task for reading sidedata.
 */
public class ReadTask<K extends Packable, V extends Packable> extends Task implements InputDrivenTask<List<byte[]>> {
  private static Logger logger = SssServer.getLogger(ReadTask.class);

  private AtomicInteger readLatch;
  private final TupleGroupScaner<K, V> sideDataScan;
  private final Queue<StorageNode> nodes;
  private Throwable exception = null;
  private boolean done = false;;
  private Class<K> keyClass;
  private Class<V> valClass;

  public ReadTask(TaskSet taskSet,
                  UUID jobID,
                  SssServerResource sssResource,
                  Configuration conf,
                  StorageNode storageNode,
                  GroupID input,
                  Class<K> keyClass,
                  Class<V> valClass,
                  TupleGroupScaner<K, V> listener) throws SssException {
    super(taskSet, jobID, sssResource, conf, storageNode, new Input(input, keyClass, valClass));
    this.sideDataScan = listener;
    this.nodes = new LinkedList<StorageNode>();
    this.keyClass = keyClass;
    this.valClass = valClass;
    this.readLatch = new AtomicInteger(input.getNSlot());
    if (input.isBroadcast()) {
      nodes.add(storageNode);
    } else {
      nodes.addAll(sssResource.getStorageNodeList());
      int n = nodes.size();
      while (!nodes.peek().equals(storageNode)) {
        nodes.offer(nodes.poll());
        n--;
        assert n >= 0;
      }
    }
  }

  public synchronized void execute() throws SssException {
    submit();
    while (!done) {
      try {
        wait();
      }
      catch (InterruptedException e) {
        // Ignore.
      }
    }
    if (exception != null) {
      throw new SssException(exception);
    }
  }

  private void submit() {
    StorageNode s = nodes.poll();
    if (s != null) {
      try {
        readLatch.set(getInputGID().getNSlot());
        submitTaskSlice(s);
      } catch (Throwable e) {
        errorCallback(e);
      }
    } else {
      synchronized (this) {
        this.done = true;
        this.notifyAll();
      }
    }
  }

  public void submitTaskSlice(StorageNode node) throws SssException {
    for (int i = 0; i < getInputGID().getNSlot(); i++) {
      DataReader r = getIOManager().createDataReader(
          getConfiguration(), node, getInputGID(), i,
          new TaskSliceSubmitter<List<byte[]>>(getWKTP(), this));
      submitReadRequest(r);
    }
  }

  @Override
  public void handle(List<byte[]> kvs) {
    try {
      Decoder keyDecoder = getKeyDecoder();
      Decoder valDecoder = getValueDecoder();

      if (getSssServerConfiguration().ioConf.stream_protocol) {
        for (int i = 0; i < kvs.size(); i += 2) {
          sideDataScan.set(
              keyClass.cast(keyDecoder.decode(kvs.get(i))),
              valClass.cast(valDecoder.decode(kvs.get(i + 1))));
        }
      }
      else {
        KeyFormatter kf = KeyFormatter.get(getInputGID().getKeyFormat());
        for (int i = 0; i < kvs.size(); i += 2) {
          sideDataScan.set(
              keyClass.cast(keyDecoder.decode(kf.getUserKey(kvs.get(i)))),
              valClass.cast(valDecoder.decode(kvs.get(i + 1))));
        }
      }
    } catch (Throwable e) {
      errorCallback(e);
    }
  }

  @Override
  public synchronized void noMoreTaskSlice(Throwable exp) {
    if (exp == null) {
      logger.debug("noMoreIORequestRead {}.", readLatch.get());
      if (readLatch.decrementAndGet() > 0) {
        return;
      }
      submit();
    } else {
      errorCallback(exp);
    }
  }

  private synchronized void errorCallback(Throwable e) {
    if (exception == null) {
      exception = e;
      done = true;
      notifyAll();
    }
  }
}
